{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from sympy import symbols, cos, sin, pi, simplify, pprint, tan, expand_trig, sqrt, trigsimp, atan2\n",
    "from sympy.matrices import Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rotx(q):\n",
    "\n",
    "  sq, cq = sin(q), cos(q)\n",
    "\n",
    "  r = Matrix([\n",
    "    [1., 0., 0.],\n",
    "    [0., cq,-sq],\n",
    "    [0., sq, cq]\n",
    "  ])\n",
    "    \n",
    "  return r\n",
    "\n",
    "\n",
    "def roty(q):\n",
    "\n",
    "  sq, cq = sin(q), cos(q)\n",
    "\n",
    "  r = Matrix([\n",
    "    [ cq, 0., sq],\n",
    "    [ 0., 1., 0.],\n",
    "    [-sq, 0., cq]\n",
    "  ])\n",
    "    \n",
    "  return r\n",
    "\n",
    "\n",
    "def rotz(q):\n",
    "\n",
    "  sq, cq = sin(q), cos(q)\n",
    "\n",
    "  r = Matrix([\n",
    "    [cq,-sq, 0.],\n",
    "    [sq, cq, 0.],\n",
    "    [0., 0., 1.]\n",
    "  ])\n",
    "    \n",
    "  return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha, beta, gamma = symbols('alpha beta gamma', real = True)\n",
    "r = rotz(alpha) * roty(beta) * rotx(gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⎡1.0⋅cos(α)⋅cos(β)  -1.0⋅sin(α)⋅cos(γ) + sin(β)⋅sin(γ)⋅cos(α)  1.0⋅sin(α)⋅sin(\n",
      "⎢                                                                             \n",
      "⎢1.0⋅sin(α)⋅cos(β)  sin(α)⋅sin(β)⋅sin(γ) + 1.0⋅cos(α)⋅cos(γ)   sin(α)⋅sin(β)⋅c\n",
      "⎢                                                                             \n",
      "⎣   -1.0⋅sin(β)                 1.0⋅sin(γ)⋅cos(β)                         1.0⋅\n",
      "\n",
      "γ) + sin(β)⋅cos(α)⋅cos(γ)⎤\n",
      "                         ⎥\n",
      "os(γ) - 1.0⋅sin(γ)⋅cos(α)⎥\n",
      "                         ⎥\n",
      "cos(β)⋅cos(γ)            ⎦\n"
     ]
    }
   ],
   "source": [
    "pprint(simplify(r))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "R00\n",
      "-1.0⋅sin(α)⋅cos(γ) + sin(β)⋅sin(γ)⋅cos(α)\n",
      "\n",
      "\n",
      "R01\n",
      "-1.0⋅sin(α)⋅cos(γ) + sin(β)⋅sin(γ)⋅cos(α)\n",
      "\n",
      "\n",
      "R02\n",
      "-1.0⋅sin(α)⋅cos(γ) + sin(β)⋅sin(γ)⋅cos(α)\n",
      "\n",
      "\n",
      "R10\n",
      "-1.0⋅sin(α)⋅cos(γ) + sin(β)⋅sin(γ)⋅cos(α)\n",
      "\n",
      "\n",
      "R11\n",
      "-1.0⋅sin(α)⋅cos(γ) + sin(β)⋅sin(γ)⋅cos(α)\n",
      "\n",
      "\n",
      "R12\n",
      "-1.0⋅sin(α)⋅cos(γ) + sin(β)⋅sin(γ)⋅cos(α)\n",
      "\n",
      "\n",
      "R20\n",
      "-1.0⋅sin(α)⋅cos(γ) + sin(β)⋅sin(γ)⋅cos(α)\n",
      "\n",
      "\n",
      "R21\n",
      "-1.0⋅sin(α)⋅cos(γ) + sin(β)⋅sin(γ)⋅cos(α)\n",
      "\n",
      "\n",
      "R22\n",
      "-1.0⋅sin(α)⋅cos(γ) + sin(β)⋅sin(γ)⋅cos(α)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for m in [0, 1, 2]:\n",
    "  for n in [0, 1, 2]: \n",
    "    \n",
    "    print()\n",
    "    print(\"R\" + str(m) + str(n))\n",
    "    pprint(r[0, 1])\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "1.0⋅sin(β)\n",
      "──────────\n",
      " │cos(β)│ \n"
     ]
    }
   ],
   "source": [
    "# beta\n",
    "sin_beta = -r[2, 0]\n",
    "cos_beta = simplify(sqrt(r[0, 0]**2 + r[1, 0]**2))\n",
    "\n",
    "print()\n",
    "pprint(sin_beta / cos_beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0⋅sin(γ)\n",
      "──────────\n",
      "  cos(γ)  \n"
     ]
    }
   ],
   "source": [
    "# gamma\n",
    "sin_gamma = r[2, 1]\n",
    "cos_gamma = r[2, 2]\n",
    "\n",
    "pprint(sin_gamma / cos_gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0⋅sin(α)\n",
      "──────────\n",
      "  cos(α)  \n"
     ]
    }
   ],
   "source": [
    "# alpha \n",
    "sin_alpha = r[1, 0]\n",
    "cos_alpha = r[0, 0]\n",
    "\n",
    "pprint(sin_alpha / cos_alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "q4, q5, q6 = symbols('q4:7', real = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⎡-sin(q₄)⋅sin(q₆) + cos(q₄)⋅cos(q₅)⋅cos(q₆)  -sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄\n",
      "⎢                                                                             \n",
      "⎢             sin(q₅)⋅cos(q₆)                             -sin(q₅)⋅sin(q₆)    \n",
      "⎢                                                                             \n",
      "⎣-sin(q₄)⋅cos(q₅)⋅cos(q₆) - sin(q₆)⋅cos(q₄)  sin(q₄)⋅sin(q₆)⋅cos(q₅) - cos(q₄)\n",
      "\n",
      ")⋅cos(q₅)  -sin(q₅)⋅cos(q₄)⎤\n",
      "                           ⎥\n",
      "               cos(q₅)     ⎥\n",
      "                           ⎥\n",
      "⋅cos(q₆)   sin(q₄)⋅sin(q₅) ⎦\n"
     ]
    }
   ],
   "source": [
    "R36 = Matrix([\n",
    "  [-sin(q4)*sin(q6) + cos(q4)*cos(q5)*cos(q6), -sin(q4)*cos(q6) - sin(q6)*cos(q4)*cos(q5), -sin(q5)*cos(q4)],\n",
    "  [                           sin(q5)*cos(q6),                           -sin(q5)*sin(q6),          cos(q5)],\n",
    "  [-sin(q4)*cos(q5)*cos(q6) - sin(q6)*cos(q4),  sin(q4)*sin(q6)*cos(q5) - cos(q4)*cos(q6),  sin(q4)*sin(q5)]])\n",
    "\n",
    "pprint(R36)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "R00\n",
      "-sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄)⋅cos(q₅)\n",
      "\n",
      "\n",
      "R01\n",
      "-sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄)⋅cos(q₅)\n",
      "\n",
      "\n",
      "R02\n",
      "-sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄)⋅cos(q₅)\n",
      "\n",
      "\n",
      "R10\n",
      "-sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄)⋅cos(q₅)\n",
      "\n",
      "\n",
      "R11\n",
      "-sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄)⋅cos(q₅)\n",
      "\n",
      "\n",
      "R12\n",
      "-sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄)⋅cos(q₅)\n",
      "\n",
      "\n",
      "R20\n",
      "-sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄)⋅cos(q₅)\n",
      "\n",
      "\n",
      "R21\n",
      "-sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄)⋅cos(q₅)\n",
      "\n",
      "\n",
      "R22\n",
      "-sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄)⋅cos(q₅)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for m in [0, 1, 2]:\n",
    "  for n in [0, 1, 2]: \n",
    "    \n",
    "    print()\n",
    "    print(\"R\" + str(m) + str(n))\n",
    "    pprint(R36[0, 1])\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sin(q₄)\n",
      "───────\n",
      "cos(q₄)\n"
     ]
    }
   ],
   "source": [
    "#q4 \n",
    "\n",
    "sin_q4 = R36[2, 2]\n",
    "cos_q4 =  -R36[0, 2]\n",
    "pprint(sin_q4 / cos_q4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "│sin(q₅)│\n",
      "─────────\n",
      " cos(q₅) \n"
     ]
    }
   ],
   "source": [
    "#q5 \n",
    "\n",
    "sin_q5 = sqrt(R36[0, 2]**2 + R36[2, 2]**2) \n",
    "cos_q5 = R36[1, 2]\n",
    "pprint(trigsimp(sin_q5 / cos_q5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sin(q₆)\n",
      "───────\n",
      "cos(q₆)\n"
     ]
    }
   ],
   "source": [
    "#q6\n",
    "\n",
    "sin_q6 = -R36[1, 1]\n",
    "cos_q6 = R36[1, 0] \n",
    "pprint(sin_q6 / cos_q6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_last_three_angles(R):\n",
    "    \n",
    "  sin_q4 = R[2, 2]\n",
    "  cos_q4 =  -R[0, 2]\n",
    "    \n",
    "  sin_q5 = sqrt(R[0, 2]**2 + R[2, 2]**2) \n",
    "  cos_q5 = R[1, 2]\n",
    "    \n",
    "  sin_q6 = -R[1, 1]\n",
    "  cos_q6 = R[1, 0] \n",
    "  \n",
    "  q4 = atan2(sin_q4, cos_q4)\n",
    "  q5 = atan2(sin_q5, cos_q5)\n",
    "  q6 = atan2(sin_q6, cos_q6)\n",
    "    \n",
    "  return q4, q5, q6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "q4:  1.6330000000000002\n",
      "\n",
      " q4: -1.50712396456497 + pi \n",
      " q5: 1.52050793847535 \n",
      " q6: -0.815787839019348\n",
      "q4: 1.63446868902482\n"
     ]
    }
   ],
   "source": [
    "R = Matrix([\n",
    "    [0.724598074610823,  -0.686235503719083,   0.0635489079820028], \n",
    "    [0.684428504026748,   0.727345057020894,   0.0502671950976013], \n",
    "    [-0.0807171180481332, 0.00707117123884714, 0.996711967115533]\n",
    "])\n",
    "\n",
    "print(\"q4: \", -1.507 + 3.14)\n",
    "# q4 = 1.6344652724032285, q5 =1.5205000259943031, q6 = -0.81578130619968381\n",
    "\n",
    "q4, q5, q6 = get_last_three_angles(R)\n",
    "print(\"\\n q4:\", q4, \"\\n q5:\", q5, \"\\n q6:\", q6)\n",
    "\n",
    "print(\"q4:\", q4.evalf())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
